# pip install mem0ai

import os
import json
from mem0 import Memory
from mem0.configs.base import MemoryConfig, LlmConfig, EmbedderConfig

# 创建配置
config = {
    "llm": {
        "provider": "ollama",
        "config": {
            "model": "qwen3:8b",
            "temperature": 0.1,
            "max_tokens": 2000,
            "ollama_base_url": "http://localhost:11434",
        }
    },
    "embedder": {
        "provider": "ollama",
        "config": {
            "model": "bge-m3:latest",
            "ollama_base_url": "http://localhost:11434",
        }
    },
    "vector_store": {
        "provider": "chroma",
        "config": {
            "collection_name": "test",
            "path": "./chroma_db",
        }
    },
    # "reset_vector_store": True
}

# 创建 Memory 实例
m = Memory.from_config(config)

# 添加对话记忆
messages = [
    {"role": "user", "content": "我喜欢科幻电影，尤其是星际穿越。"},
    {"role": "assistant", "content": "星际穿越确实是一部经典的科幻电影！我会记住你喜欢这类电影。"},
    {"role": "user", "content": "我也喜欢克里斯托弗·诺兰的其他电影。"},
    {"role": "assistant", "content": "诺兰的作品确实很出色！除了《星际穿越》，他还导演了《盗梦空间》、《信条》、《记忆碎片》等经典作品。您对这些电影有什么特别的看法吗？"}
]

# 添加记忆，并附加元数据
m.add(messages, user_id="alice", metadata={"category": "movies"})

# 检索记忆
memories = m.search(query="我喜欢什么类型的电影？", user_id="alice")
# print(memories)
print(json.dumps(memories, indent=2, ensure_ascii=False))

print("------------------------------")
# 获取所有记忆
all_memories = m.get_all(user_id="alice")
# print(all_memories)
print(json.dumps(all_memories, indent=2, ensure_ascii=False))
